module frege.data.Or where

-- adapted from http://hackage.haskell.org/packages/archive/data-or/1.0.0/doc/html/Data-Or.html
-- with some new functions and instances

import frege.control.Monoid

data Or a b = Fst a | Both a b | Snd b

--derive Read Or
derive Show Or a b 
derive Eq Or a b 

-- | Functional eliminator for 'Or'.
elimOr :: (a -> c) -> (a -> b -> c) -> (b -> c) -> Or a b -> c
elimOr f _ _ (Fst  a)   = f a
elimOr _ g _ (Both a b) = g a b
elimOr _ _ h (Snd    b) = h   b

-- | Convert an 'Either' into an 'Or'.
eitherOr :: Either a b -> Or a b
eitherOr (Left  a) = Fst a
eitherOr (Right b) = Snd b

-- convert pair of maybes into 'Or' (new, partial function)
maybeOr :: Maybe a -> Maybe b -> Or a b
maybeOr (Just a) Nothing = Fst a
maybeOr (Just a) (Just b) = Both a b
maybeOr Nothing (Just b) = Snd b
maybeOr Nothing Nothing = error "Or can't be empty" 

-- extractor for first argument (new)
orFstMaybe :: Or a b -> Maybe a
orFstMaybe (Fst a) = Just a
orFstMaybe (Both a _) = Just a
orFstMaybe _ = Nothing

-- extractor for first argument with default (new)
orFst :: Or a b -> a -> a
orFst (Fst a) _ = a
orFst (Both a _) _ = a
orFst _  a = a

-- extractor for second argument (new)
orSndMaybe :: Or a b -> Maybe b
orSndMaybe (Snd b) = Just b
orSndMaybe (Both _ b) = Just b
orSndMaybe _ = Nothing

-- extractor for second argument with default (new)
orSnd :: Or a b -> b -> b
orSnd (Snd b) _ = b
orSnd (Both _ b) _ = b
orSnd _ b = b

-- extractor for both arguments (new)
orPairMaybe :: Or a b -> (Maybe a, Maybe b)
orPairMaybe x = (orFstMaybe x, orSndMaybe x)

-- extractor for both arguments with defaults (new)
orPair :: Or a b -> a -> b -> (a, b)
orPair o a b = (orFst o a, orSnd o b)

-- | A variant of 'zip' which exhausts both lists, annotating which
-- list the elements came from. It will return zero or more @Both@,
-- followed by either zero or more @Fst@ or else zero or more @Snd@.
zipOr :: [a] -> [b] -> [Or a b]
zipOr xs ys = zipOrWithBy id (:) [] xs ys


-- | A variant of 'zipOr' with a custom 'Or'-homomorphism.
zipOrWith :: (Or a b -> c) -> [a] -> [b] -> [c]
zipOrWith k xs ys = zipOrWithBy k (:) [] xs ys

-- | A variant of 'zipOr' with a custom list-homomorphism.
zipOrBy :: (Or a b -> c -> c) -> c -> [a] -> [b] -> c
zipOrBy f z xs ys = zipOrWithBy id f z xs ys


-- | A variant of 'zipOr' with both a custom 'Or'-homomorphism and
-- a custom list-homomorphism. This is no more powerful than
-- 'zipOrBy', but it may be more convenient to separate the handling
-- of 'Or' from the handling of @(:)@.
zipOrWithBy
    :: (Or a b -> c)    -- ^ 'Or' homomorphism
    -> (c -> d -> d)    -- ^ list homomorphism, @(:)@ part
    -> d                -- ^ list homomorphism, @[]@ part
    -> [a] -> [b] -> d
zipOrWithBy k f z xs ys = go xs ys where
    go []     []     = z
    go []     (y:ys) = f (k (Snd    y)) (go [] ys)
    go (x:xs) []     = f (k (Fst  x  )) (go xs [])
    go (x:xs) (y:ys) = f (k (Both x y)) (go xs ys)

-- (new)
instance Functor Or a where
   fmap f (Snd b) = Snd (f b)
   fmap f (Both a b) = Both a (f b)
   fmap _ (Fst a) = Fst a
   
-- (new)
instance Monad Semigroup a => Or a where
   return b = Snd b
   Fst a >>= _ = Fst a
   Snd b >>= f = f b
   Both a b >>= f = case f b of
                    Fst a' -> Fst (a `mappend` a')    
                    Snd b' -> Both a b'
                    Both a' b' -> Both (a `mappend` a') b' 

-- (new)   
instance Semigroup (Semigroup a, Semigroup b) => Or a b where
   mappend x y = maybeOr (orFstMaybe x `mappend` orFstMaybe y)
                         (orSndMaybe x `mappend` orSndMaybe y)

  
